#!/usr/bin/env python
import tensorflow as tf
tf.compat.v1.disable_v2_behavior()
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
# 简单的计算来测试 GPU
# tf.debugging.set_log_device_placement(True)

from model import OpenNsfwModel, InputType
from image_utils import create_tensorflow_image_loader
from image_utils import create_yahoo_image_loader
from image_utils import create_vidio_image_loader
import numpy as np
import time
import cv2

from tensorflow.python.saved_model import builder,signature_constants,signature_def_utils,tag_constants



IMAGE_LOADER_TENSORFLOW = "tensorflow"
IMAGE_LOADER_YAHOO = "yahoo"



class NSFW:

    def __init__(self):
        self.fn_load_image = None
        self.sess = None
        self.model = OpenNsfwModel()
        self.session_run()

    def session_run(self):
        self.sess = tf.compat.v1.Session()
        input_type = InputType[InputType.TENSOR.name.lower().upper()]
        self.model.build(
            weights_path='data/open_nsfw-weights.npy', input_type=input_type)

        if input_type == InputType.TENSOR:
            if IMAGE_LOADER_YAHOO == IMAGE_LOADER_TENSORFLOW:
                self.fn_load_image = create_tensorflow_image_loader(
                    tf.Session(graph=tf.Graph()))
            else:
                # self.fn_load_image = create_yahoo_image_loader()
                self.fn_load_image = create_vidio_image_loader()
        elif input_type == InputType.BASE64_JPEG:
            import base64
            self.fn_load_image = lambda filename: np.array(
                [base64.urlsafe_b64encode(open(filename, "rb").read())])
        self.sess.run(tf.compat.v1.global_variables_initializer())


    def predict(self, input_file):
        image = self.fn_load_image(input_file)
        predictions = self.sess.run(self.model.predictions, feed_dict={
                                    self.model.input: image})
        return predictions
        #print("Results for '{}'".format(input_file))
        #print("\tSFW score:\t{}\n\tNSFW score:\t{}".format(*predictions[0]))
        # print(float(predictions[0][1]), type(float(predictions[0][1])))
        # if float(predictions[0][1]) >= 0.8:
        #     print("-------------")
    def savedModelFormat(self,export_dir):
        tf_builder = builder.SavedModelBuilder(export_dir)
        inputs = {'input':self.model.input_tensor}
        ouputs = {'output':self.model.predictions}
        signature_def_map = {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:signature_def_utils.predict_signature_def(inputs,ouputs)}
        tf_builder.add_meta_graph_and_variables(self.sess,[tf.saved_model.SERVING],strip_default_attrs=True,signature_def_map=signature_def_map)
        tf_builder.save()
    def predictVidio(self,vidio_path):
        sample_rate = 1  # 每隔100帧采样一次
        # 打开视频文件
        cap = cv2.VideoCapture(vidio_path)
        frame_count = 0
        batch = 20
        images = []
        totalNum = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame_count += 1
            if frame_count % sample_rate == 0:
                images.append(frame)
                if len(images) < batch:
                    continue
                # 处理采样帧
                images_input = self.fn_load_image(images)
                predictions = self.sess.run(self.model.predictions, feed_dict={
                    self.model.input: images_input})
                print(predictions)
                totalNum += len(images)
                images = []
        if len(images) > 0:
            images_input = self.fn_load_image(images)
            predictions = self.sess.run(self.model.predictions, feed_dict={
                self.model.input: images_input})
            print(predictions)
            totalNum += len(images)
        print("frame num:{}".format(totalNum))


if __name__ == "__main__":
    nsfw = NSFW()
    nsfw.savedModelFormat("./data/savedModelFormat")
    # total_start_time = time.time()
    # data_result = nsfw.predictVidio("./img/big-buck-bunny_trailer.mp4")
    # print(time.time()-total_start_time)
